# main.py
import os
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats
from tqdm import tqdm

from osc04a.static.constants import PlotConstants

warnings.filterwarnings('ignore')

# ---------- PERT Theoretical Distribution ----------
class PertDistribution:
    @staticmethod
    def rvs(a, b, c, size=None, random_state=None):
        alpha, beta = PertDistribution.alpha_beta(a, b, c)
        return a + (c - a) * stats.beta.rvs(alpha, beta, size=size, random_state=random_state)

    @staticmethod
    def pdf(x, a, b, c):
        alpha, beta = PertDistribution.alpha_beta(a, b, c)
        return stats.beta.pdf((x - a) / (c - a), alpha, beta) / (c - a)

    @staticmethod
    def theoretical_skewness(a, b, c):
        alpha, beta = PertDistribution.alpha_beta(a, b, c)
        return float(stats.beta(alpha, beta).stats(moments='s'))

    @staticmethod
    def theoretical_kurtosis(a, b, c):
        alpha, beta = PertDistribution.alpha_beta(a, b, c)
        return float(stats.beta(alpha, beta).stats(moments='k'))

    @staticmethod
    def theoretical_mean(a, b, c):
        return (a + 4 * b + c) / 6

    @staticmethod
    def theoretical_var(a, b, c):
        alpha, beta = PertDistribution.alpha_beta(a, b, c)
        scale = c - a
        return scale ** 2 * alpha * beta / ((alpha + beta) ** 2 * (alpha + beta + 1))

    @staticmethod
    def alpha_beta(a, b, c):
        alpha = 1 + 4 * (b - a) / (c - a)
        beta = 1 + 4 * (c - b) / (c - a)
        return alpha, beta


# ---------- 主流程 ----------
class SamplingProcessor:
    def __init__(self, config_path):
        self.cfg = pd.read_csv(config_path)
        self.res = []

    # ---------------- 工具函数 ----------------
    @staticmethod
    def _metrics(samples):
        return dict(mean=np.mean(samples),
                    var=np.var(samples),
                    median=np.median(samples),
                    skew=stats.skew(samples),
                    kurt=stats.kurtosis(samples, fisher=True))

    def _ks(self, samples, a, b, c):
        theory = PertDistribution.rvs(a, b, c, size=10000, random_state=42)
        _, p = stats.ks_2samp(samples, theory)
        return p >= 0.05, p

    # ---------------- 预采样 ----------------
    def pre_sampling(self, row):
        a, b, c = row['a'], row['b'], row['c']
        n0, inc, n_max = row['initial_n'], row['increment_n'], row['max_total_n']
        seed = row['random_seed']
        th_sk, th_ku = row['skewness_threshold'], row['kurtosis_threshold']

        n = n0
        while n <= n_max:
            s = PertDistribution.rvs(a, b, c, size=n, random_state=seed)
            m = self._metrics(s)
            if abs(m['skew'] - PertDistribution.theoretical_skewness(a, b, c)) > th_sk or \
               abs(m['kurt'] - PertDistribution.theoretical_kurtosis(a, b, c)) > th_ku:
                n += inc
                seed += 1
                continue

            if n > n0:
                s_prev = PertDistribution.rvs(a, b, c, size=n - inc, random_state=seed)
                m_prev = self._metrics(s_prev)
                if abs(m['mean'] - m_prev['mean']) / abs(m_prev['mean']) > 0.005 or \
                   abs(m['var'] - m_prev['var']) / abs(m_prev['var']) > 0.03:
                    n += inc
                    seed += 1
                    continue
            return n, s, seed
        return n - inc, s, seed

    # ---------------- 正式采样 ----------------
    def formal_sampling(self, row, n_unified, seed):
        a, b, c = row['a'], row['b'], row['c']
        th_sk, th_ku = row['skewness_threshold'], row['kurtosis_threshold']
        for _ in range(100):
            s = PertDistribution.rvs(a, b, c, size=n_unified, random_state=seed)
            m = self._metrics(s)
            if abs(m['skew'] - PertDistribution.theoretical_skewness(a, b, c)) > th_sk or \
               abs(m['kurt'] - PertDistribution.theoretical_kurtosis(a, b, c)) > th_ku:
                seed += 1
                continue
            ok, p = self._ks(s, a, b, c)
            return s, ok, p, True, seed
        return s, False, None, False, seed

    # ---------------- 绘图 ----------------
    def generate_plots(self, samples, floor, out_path):
        row = self.cfg[self.cfg['floor'] == floor].iloc[0]
        a, b, c = row['a'], row['b'], row['c']
        m = self._metrics(samples)

        alpha, beta = PertDistribution.alpha_beta(a, b, c)
        rv = stats.beta(alpha, beta, loc=a, scale=c - a)

        fig, axes = plt.subplots(2, 2, figsize=(12, 9))
        fig.suptitle(f'Validation of Empirical Duration Samples for Floor {floor[-2:]} '
                     f'from PERT({a}, {b}, {c})',
                     fontsize=16, weight='bold')

        x = np.linspace(a, c, 500)

        # (a) PDF
        ax = axes[0, 0]
        ax.plot(x, rv.pdf(x), color=PlotConstants.THEO_COLOR, lw=2,
                ls=PlotConstants.THEO_LINE, label='Theoretical')
        ax.hist(samples, bins='auto', density=True, alpha=0.4,
                color=PlotConstants.EMP_COLOR, label='Empirical')
        ax.set_xlabel('Value')
        ax.set_ylabel('Density')
        ax.set_title('(a) Empirical vs Theoretical PDF')
        ax.legend()

        # (b) Q-Q
        ax = axes[0, 1]
        stats.probplot(samples, dist=rv, plot=ax)
        ax.set_xlabel('Theoretical Quantiles')
        ax.set_ylabel('Empirical Quantiles')
        ax.set_title('(b) Q-Q Plot')
        ax.get_lines()[0].set_markerfacecolor('steelblue')
        ax.get_lines()[0].set_markersize(4)

        # (c) CDF
        ax = axes[1, 0]
        ax.plot(x, rv.cdf(x), color=PlotConstants.THEO_COLOR, lw=2,
                ls=PlotConstants.THEO_LINE, label='Theoretical')
        ax.step(np.sort(samples),
                np.linspace(1 / len(samples), 1, len(samples)),
                where='post', color=PlotConstants.EMP_COLOR, lw=2,
                ls=PlotConstants.EMP_LINE, label='Empirical')
        ax.set_xlabel('Value')
        ax.set_ylabel('Cumulative Probability')
        ax.set_title('(c) Empirical vs Theoretical CDF')
        ax.legend()

        # (d) Box
        ax = axes[1, 1]
        ax.boxplot(samples, vert=True, patch_artist=True,
                   boxprops=dict(facecolor='steelblue', alpha=0.6))
        ax.set_ylabel('Value')
        ax.set_title('(d) Empirical Box Plot')

        # 注释
        text = (f'Skewness: {m["skew"]:.3f}  '
                f'(theory {PertDistribution.theoretical_skewness(a, b, c):.3f})\n'
                f'Kurtosis: {m["kurt"]:.3f}  '
                f'(theory {PertDistribution.theoretical_kurtosis(a, b, c):.3f})')
        axes[0, 0].annotate(text, xy=(0.98, 0.97), xycoords='axes fraction',
                            ha='right', va='top',
                            bbox=dict(boxstyle='round,pad=0.4', fc='white', alpha=0.8))

        # 统一标题下移
        for ax in axes.flat:
            ax.set_title(ax.get_title(), y=-0.35)
        plt.tight_layout(rect=[0, 0.05, 1, 0.96])

        os.makedirs(os.path.dirname(out_path), exist_ok=True)
        plt.savefig(out_path, dpi=300, bbox_inches='tight')
        plt.close()

    def plot_summary_bias(self, df):
        floors = df['floor']
        metrics = ['mean', 'var', 'skewness', 'kurtosis']
        labels = [m.capitalize() for m in metrics]

        fig, axes = plt.subplots(2, 2, figsize=(max(10, 0.7 * len(floors)), 7))
        axes = axes.ravel()

        for ax, m, label in zip(axes, metrics, labels):
            ax.plot(floors, df[f'empirical_{m}'], 'o-', label='Empirical')
            ax.plot(floors, df[f'theoretical_{m}'], 's--', label='Theoretical')
            ax.set_title(label)
            ax.set_xticks(floors)
            ax.set_xticklabels(floors, rotation=45, ha='right')
            ax.legend()
            ax.grid(alpha=0.3)

        fig.suptitle('Floor-wise Empirical vs Theoretical Statistics', y=0.98)
        plt.tight_layout(rect=[0, 0, 1, 0.97])

        out_dir = os.path.dirname(self.cfg.iloc[0]['samples_validation_output_path'])
        os.makedirs(out_dir, exist_ok=True)
        plt.savefig(os.path.join(out_dir, 'summary_empirical_vs_theory.png'), dpi=300)
        plt.close()

    # ---------------- 主流程 ----------------
    def run(self):
        pre = {}
        for _, row in tqdm(self.cfg.iterrows(), desc='Pre-sampling'):
            n, s, seed = self.pre_sampling(row)
            pre[row['floor']] = {'n': n, 's': s, 'seed': seed}

        n_uni = max(v['n'] for v in pre.values())
        tqdm.write(f'Unified sample size: {n_uni}')

        for _, row in tqdm(self.cfg.iterrows(), desc='Formal'):
            fl = row['floor']
            if pre[fl]['n'] == n_uni:
                s, ok, p, conv, seed = pre[fl]['s'], *self._ks(pre[fl]['s'], row['a'], row['b'], row['c']), True, pre[fl]['seed']
            else:
                s, ok, p, conv, seed = self.formal_sampling(row, n_uni, pre[fl]['seed'])

            os.makedirs(os.path.dirname(row['samples_output_path']), exist_ok=True)
            pd.DataFrame({'value': s}).to_csv(row['samples_output_path'], index=False)

            self.generate_plots(s, fl, row['samples_validation_output_path'])

            m = self._metrics(s)
            self.res.append({
                'floor': fl,
                'final_n': n_uni,
                'empirical_mean': m['mean'],
                'theoretical_mean': PertDistribution.theoretical_mean(row['a'], row['b'], row['c']),
                'empirical_var': m['var'],
                'theoretical_var': PertDistribution.theoretical_var(row['a'], row['b'], row['c']),
                'empirical_skewness': m['skew'],
                'theoretical_skewness': PertDistribution.theoretical_skewness(row['a'], row['b'], row['c']),
                'empirical_kurtosis': m['kurt'],
                'theoretical_kurtosis': PertDistribution.theoretical_kurtosis(row['a'], row['b'], row['c']),
                'ks_pvalue': p,
                'used_seed': seed,
                'converged': conv
            })

        df = pd.DataFrame(self.res)
        out_dir = os.path.dirname(self.cfg.iloc[0]['samples_validation_output_path'])
        os.makedirs(out_dir, exist_ok=True)
        df.to_csv(os.path.join(out_dir, 'samples_validation.csv'), index=False)
        self.plot_summary_bias(df)
        print('All tasks completed.')


if __name__ == '__main__':
    SamplingProcessor('sampling_parameters2.csv').run()